# 基于FastAPI实现
# 安装依赖包
# pip install "fastapi[all]"
# pip install sse-starlette
# 完整代码实现：
import random
import time
import asyncio

import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from sse_starlette.sse import EventSourceResponse

app = FastAPI()

# 跨域设置，因为测试需要前端访问，所以允许所有域访问
origins = [
    "*"
]
# origins = [
#     "http://localhost",
#     "http://localhost:8082",
# ]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

async def event_generator():
    id = 0
    event = 'message'
    count = 0
    while True:
        # time.sleep(1)  # 模拟一些工作
        await asyncio.sleep(1)  # 每5秒发送一次消息
        yield f"id: {id}\nevent: {event}\nretry: {3000}\ndata: {{\"message\": \"{count}\"}}\n\n"
        count += 1
        id += 1

@app.get('/sse')
async def sse():
    # 返回SSE流
    return StreamingResponse(event_generator(), media_type='text/event-stream')

@app.post('/sse2')
async def sse2():
    # 返回SSE流
    return StreamingResponse(event_generator(), media_type='text/event-stream')


@app.get('/stream')
async def stream(request: Request):
    def new_count():
        return random.randint(1, 100)

    async def event_generator():
        index = 0
        while True:
            index += 1
            if await request.is_disconnected():
                break
            # 测试取随机数据，每次取一个随机数
            if count := new_count():
                yield {'data': count}

            await asyncio.sleep(1)

    return EventSourceResponse(event_generator())


if __name__ == '__main__':
    uvicorn.run('sse-server-fastapi:app', port=8080, reload=True)


# 接下来，开启服务，在浏览器访问：
# http://127.0.0.1:8080/stream
# 可以观察到，随机数字间隔一秒输出。
